import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

data = pd.read_csv('E:\\code\\611\\轻量化\\飞行数据\\训练集test\\20_20201102_123315.csv', encoding='UTF-8')
target_col = 'L2-ALF'

drop_col = [col for col in data.columns if col.startswith('L0-')]
data = data.drop(columns=drop_col)

#计算皮尔逊相关系数＆皮尔曼相关系数
pearson = data.corr(method='pearson')[[target_col]].drop(target_col)
spearman = data.corr(method='spearman')[[target_col]].drop(target_col)

# #降序排列，取前n个
# order_pearson = pearson.abs().sort_values(by=target_col,ascending=False).head(15)
# order_spearman = spearman.abs().sort_values(by=target_col,ascending=False).head(15)
# #输出最后一个特征的相关系数
# threshold_pearson = order_pearson.iloc[-1]
# threshold_spearman = order_spearman.iloc[-1]
# print(f'皮尔逊阈值：{threshold_pearson}，斯皮尔曼阈值：{threshold_spearman}')
#
# target_pearson = [target_col] + order_pearson.index.tolist()
# target_spearman = [target_col] + order_spearman.index.tolist()
# data_pearson = data[target_pearson]
# data_spearman = data[target_spearman]
# output_path = 'E:\\code\\611\\轻量化\\飞行数据\\processed_test'
# data_pearson.to_csv(f'{output_path}\\data_pearson.csv',index=False,encoding='utf-8')
# data_spearman.to_csv(f'{output_path}\\data_spearman.csv',index=False,encoding='utf-8')

#取均值
combined_corr = (pearson.abs() + spearman.abs()) / 2
ordered_features = combined_corr.sort_values(by=target_col, ascending=False).head(15)
target_corr = [target_col] + ordered_features.index.tolist()
data_combine = data[target_corr]
output_path = 'E:\\code\\611\\轻量化\\飞行数据\\processed_test'
data_combine.to_csv(f'{output_path}\\data_combine.csv',index=False,encoding='utf-8')

feature_list_path = 'E:\\code\\611\\轻量化\\飞行数据\\selected_features.txt'
with open(feature_list_path, 'w', encoding='utf-8') as f:
    for feature in target_corr:
        f.write(feature + '\n')

# plt.figure(figsize=(10, 6))
# sns.heatmap(order_pearson, annot=True, cmap='YlGn', fmt='.2f', linewidths=0.5, linecolor='black')
# plt.title('pearson')
# plt.show()
# plt.figure(figsize=(10, 6))
# sns.heatmap(order_spearman, annot=True, cmap='YlGn', fmt='.2f', linewidths=0.5, linecolor='black')
# plt.title('spearman')
# plt.show()

